from jili.core.printlog import print
import rqdatac
from rqdatac import LiveMarketDataClient
from jili.tool.convert import str2datetime
import datetime
import pandas as pd
ats2rq_mkcode={
    "SH":"XSHG",
    "SZ":"XSHE",
    # "BJ":"BJSE",
    "HK":"HK",
    # "SHK":"SHK",
    # "ZHK":"ZHK",
    'CFFEX':"CFFEX",
    'CZCE':"CZCE",
    'DCE':"DCE",
    'GFEX':"GFEX",
    'INE':"INE",
    'SHFE' :"SHFE"
}
rq2ats_mkcode={v: k for k, v in ats2rq_mkcode.items()}
class tick_l1():
    def __init__(self):
        self.ticks={}
    def reset(self):
        self.ticks={}
    def dftick2tick0(self,k):
        if isinstance(k, pd.DataFrame):
            if not k.empty:
                index0 = k.index.to_list()
                ids = []
                tk = []
                for ii in index0:
                    ids.append(ii[0])
                    tk.append(ii[1])
                k["timekey"] = tk
                k["order_book_id"] = ids
                k = k.to_dict(orient="records")
                kk=[]
                prei=k[0]
                for i in k:
                    if prei["trading_date"]!=i["trading_date"]:
                        self.reset()
                    if prei["open"]!=0 and i["open"]==0:
                        self.reset()
                    t=self.xtick2tick0(i["order_book_id"], i)
                    if t is None:
                        continue
                    kk.append(t)
                    prei=i
                return kk
            else:
                return []
        else:
            return []

    def dftick2tick0_f(self,k):
        if isinstance(k, pd.DataFrame):
            if not k.empty:
                index0 = k.index.to_list()
                ids = []
                tk = []
                for ii in index0:
                    ids.append(ii[0])
                    tk.append(ii[1])
                k["timekey"] = tk
                k["order_book_id"] = ids
                k = k.to_dict(orient="records")
                kk=[]
                prei=k[0]
                for i in k:
                    if prei["trading_date"]!=i["trading_date"]:
                        self.reset()
                    if prei["open"]!=0 and i["open"]==0:
                        self.reset()
                    t=self.xtick2tick0_f(i["order_book_id"], i)
                    if t is None:
                        continue
                    kk.append(t)
                    prei=i
                return kk
            else:
                return []
        else:
            return []

    def dftick2tick2_f(self, k):
        if isinstance(k, pd.DataFrame):
            if not k.empty:
                index0 = k.index.to_list()
                ids = []
                tk = []
                for ii in index0:
                    ids.append(ii[0])
                    tk.append(ii[1])
                k["timekey"] = tk
                k["order_book_id"] = ids
                k = k.to_dict(orient="records")
                kk = []
                prei = k[0]
                for i in k:
                    if prei["trading_date"] != i["trading_date"]:
                        self.reset()
                    if prei["open"] != 0 and i["open"] == 0:
                        self.reset()
                    t = self.xtick2tick2_f(i["order_book_id"], i)
                    if t is None:
                        continue
                    kk.append(t)
                    prei = i
                return kk
            else:
                return []
        else:
            return []
    def dftick2tick1(self,k):
        if isinstance(k, pd.DataFrame):
            if not k.empty:
                index0 = k.index.to_list()
                ids = []
                tk = []
                for ii in index0:
                    ids.append(ii[0])
                    tk.append(ii[1])
                k["timekey"] = tk
                k["order_book_id"] = ids
                k = k.to_dict(orient="records")
                kk=[]
                prei=k[0]
                for i in k:
                    if prei["trading_date"]!=i["trading_date"]:
                        self.reset()
                    if prei["open"]!=0 and i["open"]==0:
                        self.reset()
                    t=self.xtick2tick1(i["order_book_id"], i)
                    kk.append(t)
                    prei=i
                return kk
            else:
                return []
        else:
            return []
    def xtick2tick0_f(self,code,tick,call=None):
        timekey=tick["timekey"]
        cum_amount = tick["total_turnover"]
        cum_volume = tick["volume"]
        r=self.ticks.get(code)
        if r is None:
            self.ticks[code] = (timekey, cum_amount, cum_volume)
            amount = cum_amount
            volume = cum_volume
        else:
            timekey0, cum_amount0, cum_volume0 = r
            if timekey0 < timekey:
                amount =  cum_amount-cum_amount0
                volume =  cum_volume-cum_volume0
                self.ticks[code] = (timekey, cum_amount, cum_volume)
            else:
                return None
        t = {}
        t["timekey"] = timekey
        t["obj"]=code
        # t["code"] = t["obj"]+"."+ t["mkcode"]
        t["lastPrice"] = tick["last"]
        if cum_volume==0:
            avgPrice=t["lastPrice"]
        else:
            avgPrice=cum_amount/cum_volume
        t["avgPrice"]=avgPrice
        t["amount"] = amount
        t["volume"] = volume
        t["askPrice"] = [tick["a1"]]
        t["bidPrice"] = [tick["b1"]]
        t["askVol"] = [tick["a1_v"]]
        t["bidVol"] = [tick["b1_v"]]
        t["pos"] = tick["open_interest"]
        if call is not None:
            call(t)
        return t
    def xtick2tick2_f(self,code,tick,call=None):
        timekey=tick["timekey"]
        cum_amount = tick["total_turnover"]
        cum_volume = tick["volume"]
        r=self.ticks.get(code)
        if r is None:
            self.ticks[code] = (timekey, cum_amount, cum_volume)
            amount = cum_amount
            volume = cum_volume
        else:
            timekey0, cum_amount0, cum_volume0 = r
            if timekey0 < timekey:
                amount =  cum_amount-cum_amount0
                volume =  cum_volume-cum_volume0
                self.ticks[code] = (timekey, cum_amount, cum_volume)
            else:
                return None
        t = {}
        t["timekey"] = timekey
        t["obj"] = code
        t["lastPrice"] = tick["last"]
        if cum_volume==0:
            avgPrice=t["lastPrice"]
        else:
            avgPrice=cum_amount/cum_volume
        t["avgPrice"]=avgPrice
        t["amount"] = amount
        t["volume"] = volume
        t["askPrice"] = [tick["a1"],tick["a2"],tick["a3"],tick["a4"],tick["a5"]]
        t["bidPrice"] = [tick["b1"],tick["b2"],tick["b3"],tick["b4"],tick["b5"]]
        t["askVol"] = [tick["a1_v"],tick["a2_v"],tick["a3_v"],tick["a4_v"],tick["a5_v"]]
        t["bidVol"] = [tick["b1_v"],tick["b2_v"],tick["b3_v"],tick["b4_v"],tick["b5_v"]]
        t["pos"] = tick["open_interest"]
        if call is not None:
            call(t)
        return t
    def xtick2tick0(self,code,tick,call=None):
        timekey=tick["timekey"]
        cum_amount = tick["total_turnover"]
        cum_volume = tick["volume"]
        r=self.ticks.get(code)
        if r is None:
            self.ticks[code] = (timekey, cum_amount, cum_volume)
            amount = cum_amount
            volume = cum_volume
        else:
            timekey0, cum_amount0, cum_volume0 = r
            if timekey0 < timekey:
                amount =  cum_amount-cum_amount0
                volume =  cum_volume-cum_volume0
                self.ticks[code] = (timekey, cum_amount, cum_volume)
            else:
                return None
        t = {}
        t["timekey"] = timekey
        t["obj"], t["mkcode"] = code2obj(code)
        t["code"] = t["obj"]+"."+ t["mkcode"]
        t["lastPrice"] = tick["last"]
        if cum_volume==0:
            avgPrice=t["lastPrice"]
        else:
            avgPrice=cum_amount/cum_volume
        t["avgPrice"]=avgPrice
        t["amount"] = amount
        t["volume"] = volume
        t["askPrice"] = [tick["a1"],tick["a2"],tick["a3"],tick["a4"],tick["a5"]]
        t["bidPrice"] = [tick["b1"],tick["b2"],tick["b3"],tick["b4"],tick["b5"]]
        t["askVol"] = [tick["a1_v"],tick["a2_v"],tick["a3_v"],tick["a4_v"],tick["a5_v"]]
        t["bidVol"] = [tick["b1_v"],tick["b2_v"],tick["b3_v"],tick["b4_v"],tick["b5_v"]]
        t["num_trades"] = tick["num_trades"]
        if call is not None:
            call(t)
        return t
    def xtick2tick2(self,code,tick,call=None):
        timekey=tick["timekey"]
        cum_amount = tick["total_turnover"]
        cum_volume = tick["volume"]
        r=self.ticks.get(code)
        if r is None:
            self.ticks[code] = (timekey, cum_amount, cum_volume)
            amount = cum_amount
            volume = cum_volume
        else:
            timekey0, cum_amount0, cum_volume0 = r
            if timekey0 < timekey:
                amount =  cum_amount-cum_amount0
                volume =  cum_volume-cum_volume0
                self.ticks[code] = (timekey, cum_amount, cum_volume)
            else:
                return None
        tick["timekey"] = timekey
        tick["obj"], tick["mkcode"] = code2obj(code)
        tick["code"] = tick["obj"] + "." + tick["mkcode"]
        tick["lastPrice"] = tick["last"]
        if cum_volume==0:
            avgPrice=tick["lastPrice"]
        else:
            avgPrice=cum_amount/cum_volume
        tick["avgPrice"]=avgPrice
        tick["cumAmount"] = cum_amount
        tick["cumVolume"] = cum_volume
        tick["amount"] = amount
        tick["volume"] = volume
        tick["askPrice"] = [tick["a1"], tick["a2"], tick["a3"], tick["a4"], tick["a5"]]#,tick["a6"], tick["a7"], tick["a8"], tick["a9"], tick["a10"]]
        tick["bidPrice"] = [tick["b1"], tick["b2"], tick["b3"], tick["b4"], tick["b5"]]
        tick["askVol"] = [tick["a1_v"], tick["a2_v"], tick["a3_v"], tick["a4_v"], tick["a5_v"]]
        tick["bidVol"] = [tick["b1_v"], tick["b2_v"], tick["b3_v"], tick["b4_v"], tick["b5_v"]]
        for f0 in ["a","b"]:
            for f1 in ["1","2","3","4","5"]:
                f=f0+f1
                del tick[f]
                del tick[f+"_v"]
        if call is not None:
            call(tick)
        return tick
    def xtick2tick1(self,code,tick,call=None):
        timekey = tick["timekey"]
        cum_amount = tick["total_turnover"]
        cum_volume = tick["volume"]
        r = self.ticks.get(code)
        if r is None:
            self.ticks[code] = (timekey, cum_amount, cum_volume)
            amount = cum_amount
            volume = cum_volume
        else:
            timekey0, cum_amount0, cum_volume0 = r
            if timekey0 < timekey:
                amount = cum_amount - cum_amount0
                volume = cum_volume - cum_volume0
                self.ticks[code] = (timekey, cum_amount, cum_volume)
            else:
                return None
        tick["timekey"] = timekey
        tick["obj"], tick["mkcode"] = code2obj(code)
        tick["code"] = tick["obj"] + "." + tick["mkcode"]
        tick["lastPrice"] = tick["last"]
        if cum_volume==0:
            avgPrice=tick["lastPrice"]
        else:
            avgPrice=cum_amount/cum_volume
        tick["avgPrice"]=avgPrice
        tick["cumAmount"] = cum_amount
        tick["cumVolume"] = cum_volume
        tick["amount"] = amount
        tick["volume"] = volume
        tick["askPrice"] = [tick["a1"], tick["a2"], tick["a3"], tick["a4"], tick["a5"]]
        tick["bidPrice"] = [tick["b1"], tick["b2"], tick["b3"], tick["b4"], tick["b5"]]
        tick["askVol"] = [tick["a1_v"], tick["a2_v"], tick["a3_v"], tick["a4_v"], tick["a5_v"]]
        tick["bidVol"] = [tick["b1_v"], tick["b2_v"], tick["b3_v"], tick["b4_v"], tick["b5_v"]]
        for f0 in ["a","b"]:
            for f1 in ["1","2","3","4","5"]:
                f=f0+f1
                del tick[f]
                del tick[f+"_v"]
        if call is not None:
            call(tick)
        return tick
def obj2code(obj,mkcode=None):
    if not obj[0].isalpha():
        if mkcode is None:
            if "." in obj:
                obj,mkcode=obj.split(".")
                if mkcode in ats2rq_mkcode.keys():
                    mkcode=ats2rq_mkcode[mkcode]
            else:
                if obj[0]>"3":
                    mkcode="XSHG"
                else:
                    mkcode="XSHE"
        else:
            if mkcode in ats2rq_mkcode.keys():
                mkcode = ats2rq_mkcode[mkcode]
        return obj+"."+mkcode
    else:
        if "." in obj:
            obj, mkcode = obj.split(".")
            return obj
        else:
            return obj
def code2obj(code):
    if "." in code:
        obj,mkcode=code.split(".")
        if mkcode in rq2ats_mkcode.keys():
            mkcode=rq2ats_mkcode[mkcode]
    else:
        obj=code
        mkcode=None
    return obj,mkcode

rq2ats_products={
    "CS":"stock",
    "INDX":"index",
    "ETF":"etf",
    "LOF":"fund",
    "Spot":"spot",
    "Repo":"repo",
    "Convertible":"convertible",
    "Future":"future",
    "Option":"option"
}
ats2rq_mkcode={v: k for k, v in rq2ats_products.items()}
periods={#60s, 300s, 900s, 1800s, 3600s
    "tick": "tick",
    "k1m":"1m",
    "k3m":"3m",
    "k5m":"5m",
    "k10m":"10m",
    "k15m":"15m",
    "k30m":"30m",
    "k60m":"60m",
    "k1d":"1d",
    "k1w":"1w"
}
periods1={v: k for k, v in periods.items()}
ats2rq_fqtype={
    "qfq":"pre",
    "hfq":"post",
    "bfq":"none"
}
def ats2rq_period(period):
    if period in periods.keys():
        return periods[period]
    else:
        return period[1:]
class pricetool():
    def __init__(self,islive=False):
        self.calls={}
        self.obj_infos = {}
        self.islive=islive
        try:
            if not rqdatac.initialized():
                rqdatac.init()
            if islive:
                self.tool= LiveMarketDataClient()
        except Exception as e:
            print(e)
    def run(self):
        if self.islive:
            self.tool.listen(handler=self.handle_msg)
    def get_objinfo(self,obj):
        if obj not in self.obj_infos.keys():
            code = obj2code(obj)
            self.obj_infos[obj] = rqdatac.instruments(code, market='cn')
        return self.obj_infos[obj]
    def get_tradeday(self,count=1,end_date=None):
        if end_date is None:
            now = datetime.datetime.now()
            today = datetime.datetime(now.year, now.month, now.day)
        else:
            if isinstance(end_date,str):
                today=str2datetime(end_date)
            else:
                today=end_date
        start=today-datetime.timedelta(days=2*count+5)
        tdays=rqdatac.get_trading_dates(start_date=start,end_date=today)
        r=[]
        # count0=count+1
        count0=count
        for i in tdays[-count0:]:
            r.append(datetime.datetime(i.year,i.month,i.day))
        return r
    def istradeday(self):
        now = datetime.datetime.now()
        today = datetime.datetime(now.year, now.month, now.day)
        r=rqdatac.get_trading_dates(start_date=today,end_date=today)
        if len(r)>0:
            return True
        else:
            return False
    def subscribe(self,obj, period='k30m', callback=None):
        if isinstance(obj,str):
            codes=[obj2code(obj)]
        else:
            codes=[]
            for i in obj:
                codes.append(obj2code(i))
        period0=ats2rq_period(period)
        for code in codes:
            if period[0]=="k":
                code0="bar_"
                code1=period0
            else:
                code0="tick_"
                code1=""
            key=code0+code+"_"+code1
            if self.islive:
                self.tool.subscribe(key)
            if code not in self.calls.keys():
                self.calls[code]= {}
            self.calls[code][period]=callback



    def get_today_bar(self,obj,period="k1d"):
        bars=self.get_prebars(obj,count=1,period=period,fqtype="bfq")
        if bars:
            return bars[-1]
        else:
            return {}
    def get_prebars(self,obj,count=1,period="k1d",fqtype="qfq",start_date='',end_date=''):
        obj=obj2code(obj)
        fqtype0 = ats2rq_fqtype.get(fqtype,"none")

        period0=ats2rq_period(period)
        if end_date == '':
            end_date0 = self.get_tradeday(1)[0]
        else:
            end_date0 = end_date
        if count>0:
            if start_date=='':
                tdays=self.get_tradeday(count=count,end_date=end_date0)
                start_date=tdays[0]
        print("price_rq get_price",obj,start_date,end_date,end_date0,count,period,period0,fqtype,fqtype0)
        bars=rqdatac.get_price(obj, start_date=start_date, end_date=end_date0, frequency=period0, fields=None, adjust_type=fqtype0, skip_suspended =False, market='cn', expect_df=True,time_slice=None)
        if bars is None:
            return []
        if bars.empty:
            return []
        index0 = bars.index.to_list()
        # if isinstance(index0[0], tuple):
        obj = []
        ids = []
        tk = []
        for i in index0:
            obj.append(i[0].split(".")[0])
            ids.append(i[0])
            tk.append(i[1])
        bars["timekey"] = tk
        bars["obj"] = obj
        bars["order_book_id"] = ids
        bars["amount"] = bars["total_turnover"]
        return bars.to_dict("records")

if __name__ == "__main__":
    pt=pricetool()
    b=pt.get_prebars("000400.XSHE",fqtype="qfq",period="k15m",start_date="20250516")
    print("end")